"""
用root按下面方式调用
R_HOST=127.0.0.1 R_PORT=6379 M_HOST=127.0.0.1 M_PORT=27017 L_PORT=6000 S_PORT=7760 N_PORT=7683 python3 wechat.py


腾讯的限制： handle里面函数处理超过5秒会连续调用3次，而且客户端不会有信息返回。
https://github.com/offu/WeRoBot/issues/486
"""

import werobot
import datetime
import sys
import os
import time
from redis_interface import text_gen
import redis
import pymongo as pm
from langchain.llms import OpenAI
import threading
from neo4j import GraphDatabase

sys.path.insert(0, '/home/peiyp/code')
from PyCmpltrtok.common import sep, to_decode, to_int
from PyCmpltrtok.util_mongo import delete_many_by_user, get_history, enqueue, VALUE
from daemon import MONGODB_NAME
from xlangchain_gen import analyze_symptom4user
from python_nlp.xfaiss.flask_faiss_on_mongo_test import search
from python_nlp.kg.neo4j.seek_data_with_embed.xdiagnose import get_symptom, diagnose
from common import RESULT_PREFIX


"""
映射langchain需要的openai服务端口L_PORT到本机。
映射search函数需要的S_PORT到本机。
映射neo4j服务到本机N_PORT
"""
L_PORT = int(os.environ['L_PORT'])
S_PORT = int(os.environ['S_PORT'])

# 连接neo4j
print('Connecting neo4j ...')
driver = GraphDatabase.driver(**{
    "uri": f"bolt://127.0.0.1:{os.environ['N_PORT']}",
    "auth": ("neo4j", "p1983ncbdA1-ei"),
    "encrypted": False
})
print('Connected to neo4j.')

# langchain模型
model = OpenAI(
    streaming=True,
    verbose=True,
    callbacks=[],
    openai_api_key="token1",
    openai_api_base=f"http://127.0.0.1:{L_PORT}/v1",
    model_name="chatglm2-6b",
    temperature=0.0,
    openai_proxy=None,
    top_p=1.0,
)

# 连接Redis
rdb = redis.Redis(os.environ['R_HOST'], int(os.environ['R_PORT']), 0)
rdb.get('try_it')

# 连接Mongodb
mongo = pm.MongoClient(os.environ['M_HOST'], int(os.environ['M_PORT']), serverSelectionTimeoutMS=3000)
mdb = mongo[MONGODB_NAME]
get_history(mdb, 'user_xxxx', limit=1)  # try it

robot = werobot.WeRoBot(token='kx8M7eODU4HiyAmw3a')
robot.config['APP_ID'] = 'wx47577a029b258f6b'
robot.config['ENCODING_AES_KEY'] = 'LjBrnwSvqGXx89Tszleo4sMYp8yMALcTArwhPjxoZGJ'
client = robot.client

CHAR_LIMIT = 450
PROMPT_MED = '（诊断结果仅供参考，请咨询专业医师。）'
PROMPT_SHORT = '（请继续对话，或者输入逗号（,或，）来重置对话。）'
PROMPT = '（因平台限制，内容未完全展示，输入句号（.或。）来继续展示。请继续对话，或者输入逗号（,或，）来重置对话。）'
NO_MORE = '（已经没有需要继续输出的内容了。）'
GENERATING = '（内容生成中……）'
GO_ON_SET = set(['.', '。'])
RESET_SET = set([',', '，'])
WAIT_LAST_TURN = '（请等待上一轮对话完成。）'
THRESH_S = 45  # 单例锁最多等待这么多秒
DIS_THRESH = 20.0  # 搜索时距离上限
WORDS_LIMIT = 4  # 搜索后最多用这么多症状
THRESH_S_DIAG = 45  # 诊断超过这么多秒就不作数了


def list2str_with_limit(xlist, xlimit):
    """
    以xlimit个字符为限制，把列表转字符串。
    """
    assert(isinstance(xlimit, int) and xlimit > 0)
    xstr = ''
    xfirst = True
    for xelement in xlist:
        if len(xstr) + 2 + len(xelement) > xlimit:
            return xstr + '……'
        if xfirst:
            xfirst = False
        else:
            xstr += ', '
        xstr += str(xelement)
    return xstr


def encode_set(xset):
    """
    编码集合
    """
    if not xset:
        return ''
    return '|'.join(sorted(xset))


def decode_set(xstr):
    """
    解码编码后的集合。（注意返回是列表。）
    """
    if not xstr:
        return []
    return xstr.split('|')

if 0:
    """
    自定义菜单 需要 微信认证
    微信认证： 账号主体为个人，无法开通微信认证
    """
    client.create_menu({
        'button': [{
            'type': 'click',
            'name': '继续',
            'key': 'go_on'
        }]
    })


    @robot.key_click('go_on')
    def go_on(message):    
        pass
    

@robot.handler
def hello(message):
    
    sep()
    xuuid = ''
    
    # 判断类型
    xtype = message.type
    print(datetime.datetime.now(), xtype)
    if 'text' != xtype:
        return '目前仅支持文本信息！' + PROMPT_SHORT
    
    # 用户标识
    xuserid = message.source
    print('userid:', type(xuserid), xuserid)

    def do_return(xoutput, is_release=True):
        """
        闭包函数，通用返回过程。
        """
        print('OUTPUT:', xoutput)
        
        # 本轮聊天放入mongodb
        xinput4db = xinput
        if xinput4db in RESET_SET:
            xinput4db = '我们重新开始对话吧。'
        if xinput4db in GO_ON_SET:
            xinput4db = '您请继续说。'
        ts = time.time_ns()
        enqueue(mdb, 'dialog_in', xuserid, ts, {
            VALUE: xinput4db,
            'io': 'i',
        })
        xoutput4db = xoutput.replace(PROMPT, '').replace(PROMPT_MED, '').replace(PROMPT_SHORT, '')
        xoutput4db = xoutput4db.replace(NO_MORE, '我说完了。').replace(GENERATING, '让我想想……')
        ts = time.time_ns()
        enqueue(mdb, 'dialog_out', xuserid, ts, {
            VALUE: xoutput4db,
            'io': 'o',
        })
        
        # 解除单例限制
        if is_release:
            rdb.srem('users_now', xuserid)
            
        return xoutput

    # 输入文本
    xinput = message.content.strip()
    print('input:', xinput)

    # 单例限制
    if xinput not in GO_ON_SET:
        xis_in = rdb.sismember('users_now', xuserid)
        ts_now = time.time_ns()
        if xis_in:
            ts = rdb.hget('username2ts', xuserid)
            ts = to_int(ts)
            if ts and ts_now - ts < THRESH_S * 1e9:
                xoutput = WAIT_LAST_TURN + PROMPT_SHORT
                print('OUTPUT:', xoutput)
                return xoutput
        rdb.hset('username2ts', xuserid, ts_now)
        rdb.sadd('users_now', xuserid)

    def get_result_of_diagnose():
        """
        最后对诊断的结果进行判断和返回。
        """
        # 症状分析是否结束？
        user2result = rdb.hget('user2result', xuserid)
        if user2result is None:
            ts = rdb.hget('user2result_ts', xuserid)
            ts = to_int(ts)
            ts_now = time.time_ns()
            if ts and ts_now - ts < THRESH_S_DIAG * 1e9:  # 判断时间来防止诊断线程的问题造成永远输出（内容生成中……）
                return GENERATING + PROMPT
            else:
                rdb.hset('user2result', xuserid, '')  # 如果超时了，就设置成诊断没有结果
        user2result = to_decode(user2result).strip()
        if user2result:
            
            # 丢弃闲聊
            rdb.hdel('username2dynamic', xuserid)
            rdb.hdel('username2output', xuserid)
            if xuuid:
                rdb.hdel('uuid2output', xuuid)
            
            return user2result + PROMPT_MED + PROMPT_SHORT
        return None 

    def get_result():
        """
        闭包函数，拿数据
        """
        # 最后对诊断的结果进行判断和返回。
        r = get_result_of_diagnose()
        if r is not None:
            return r
               
        # 最后结果
        xoutput = rdb.hget('username2output', xuserid)
        
        # 当前偏移
        xoffset = rdb.hget('username2offset', xuserid)
        xoffset = to_int(xoffset)
        xoffset_ = xoffset
        # 分推理结束和未结束两种
        if xoutput is not None:  # 结束
            xoutput = xoutput.decode('utf8')
            
            # 取本次输出
            xoutput = xoutput[xoffset:]
            xlen_output = len(xoutput)
            
            # 分输出完成未完成两种
            if xlen_output <= CHAR_LIMIT:
                # 输出完成，清理username2dynamic以标志输出完成
                rdb.hdel('username2dynamic', xuserid)
                rdb.hdel('username2offset', xuserid)
                return xoutput + PROMPT_SHORT
            else:
                # 输出未完成，分段输出
                xres = xoutput[:CHAR_LIMIT]
                
        else:  # 未结束
            
            # 从动态响应里拿
            xres = rdb.hget('username2dynamic', xuserid)
            xres = to_decode(xres)
            
            # 限制一下长度
            xres = xres[xoffset:xoffset+CHAR_LIMIT]
            
        # 记录偏移
        xlen = len(xres)
        rdb.hset('username2offset', xuserid, xoffset_ + xlen)
        
        return xres + PROMPT
    
    if xinput in GO_ON_SET:  # "继续显示"
        # 最后对诊断的结果进行判断和返回。
        r = get_result_of_diagnose()
        if r is not None:
            return do_return(r)
        
        # 从动态响应里拿
        xres = rdb.hget('username2dynamic', xuserid)
        # 没有动态响应，说明输出结束
        if xres is None:
            return do_return(NO_MORE + PROMPT_SHORT)
        
        print('GO ON >>>>')
        return do_return(get_result(), is_release=False)
    elif xinput in RESET_SET:  # “重置对话”
        delete_many_by_user(mdb, 'dialog_in', xuserid)
        delete_many_by_user(mdb, 'dialog_out', xuserid)
        rdb.hdel('user2dis', xuserid)
        return do_return('对话已经重置。')
    
    def analyze_disease():
        # 以删除来标志开始
        rdb.hdel('user2result', xuserid)
        ts = time.time_ns()
        rdb.hset('user2result_ts', xuserid, ts)
        
        # 当前用户的结果
        xencoded = rdb.hget('user2dis', xuserid)
        xencoded = to_decode(xencoded)
        xdis_set = set(decode_set(xencoded))
        xdis_list_ = sorted(xdis_set)
        sep()
        print(f'当前诊断：{xdis_list_}')
    
        # 从闲聊中分析用户不适
        xsymptoms = analyze_symptom4user(mdb, xinput, xuserid, model)
        print('SYMPTOMS:', xsymptoms)
        
        # 搜索输入，选出距离小于阈值的结果
        for xsym in xsymptoms:
            sep(xsym)
            xoutput = search(xsym, 8, S_PORT)
            xsents = xoutput['sentences']
            xds = xoutput['D']
            print('候选:')
            xwords = []
            for i, s in enumerate(xsents):
                d = xds[i]
                print(i, s, d, d < DIS_THRESH)
                if d < DIS_THRESH:
                    xwords.append(s)

            # 从这些结果中挑选出是“症状”的节点名字
            xlist = get_symptom(driver, xwords, WORDS_LIMIT)
            print(f'输入了症状：{xlist}')
        
            # 用这些症状去诊断
            xdiseases = diagnose(driver, xlist)
            if xdiseases:
                print(f'本轮诊断：{xdiseases}')
                if not xdis_set:
                    # 第一次诊断
                    xdis_set = set(xdiseases)
                else:
                    # 后续诊断
                    xresult = xdis_set.intersection(set(xdiseases))  # 求交集
                    if len(xresult):
                        print('【有交集】')
                    else:
                        # 没有交集求并集
                        print('【没有交集求并集】')
                        xresult = xdis_set.union(set(xdiseases))
                    xdis_set = xresult
                    print('交并后:', xdis_set)
                    
        # 结果            
        xdis_list = sorted(xdis_set)
        
        # 结果还是那样，就不使用
        if xdis_list_ == xdis_list:
            print('【诊断结果没有变化】')
            rdb.hset('user2result', xuserid, '')
            return
        
        # 保存用户结果
        xencoded = encode_set(xdis_set)
        rdb.hset('user2dis', xuserid, xencoded.encode('utf8'))
        
        # 返回诊断
        xoutput = RESULT_PREFIX + list2str_with_limit(xdis_list, CHAR_LIMIT)
        if len(xdis_list) > 5:
            xoutput += '（可能性较多，请进一步描述症状来缩小范围。）'
        rdb.hset('user2result', xuserid, xoutput)
        
    th = threading.Thread(target=analyze_disease)
    th.start()
    
    # 生成
    xuuid = text_gen(rdb, xinput, xuserid)
    
    def check_quick_reply():
        """
        闭包函数，侦测快速返回，并获取。
        """
        # 症状分析是否结束？
        user2result = rdb.hget('user2result', xuserid)
        if user2result is None:
            return None
        user2result = user2result.decode('utf8').strip()
        if user2result:
            
            # 丢弃闲聊
            rdb.hdel('username2dynamic', xuserid)
            rdb.hdel('username2output', xuserid)
            if xuuid:
                rdb.hdel('uuid2output', xuuid)
                
            return user2result + PROMPT_MED + PROMPT_SHORT
        
        # 最后结果
        xoutput = rdb.hget('username2output', xuserid)
        
        # 如果还没有结束
        if xoutput is None:
            return None
        
        # 已经结束
        xoutput = xoutput.decode('utf8')
        
        # 当前偏移
        xoffset = rdb.hget('username2offset', xuserid)
        xoffset = to_int(xoffset)
        xoffset_ = xoffset
        
        # 取本次输出
        xoutput = xoutput[xoffset:]
        xlen_output = len(xoutput)
        
        # 分输出完成未完成两种
        if xlen_output <= CHAR_LIMIT:
            # 输出完成，清理username2dynamic以标志输出完成
            rdb.hdel('username2dynamic', xuserid)
            rdb.hdel('username2offset', xuserid)
            return xoutput + PROMPT_SHORT
        else:
            # 输出未完成，分段输出
            xres = xoutput[:CHAR_LIMIT]
            
            # 记录偏移
            xlen = len(xres)
            rdb.hset('username2offset', xuserid, xoffset_ + xlen)
            
            return xres + PROMPT
    
    # 既然限制5秒，干脆等3.75秒，让大模型去生成
    for _ in range(3):
        time.sleep(1.0)
        r = check_quick_reply()
        if r is not None:
            return do_return(r)
    time.sleep(0.75)
    
    # 3.75秒后，拿结果
    xres = get_result()
    return do_return(xres)


# 让服务器监听在 0.0.0.0:80
robot.config['HOST'] = '0.0.0.0'
robot.config['PORT'] = 80
robot.run()
